1521D - Nastia Plays with a Tree - CodeForces Solution


constructive algorithms data structures dfs and similar dp dsu greedy implementation trees *2500

Please click on ads to support us..

C++ Code:

#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, b) for(int i = a; i < (b); ++i)
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()
#define F first
#define S second
#define pb push_back
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;

int main() {
	cin.tie(0)->sync_with_stdio(0);
	cin.exceptions(cin.failbit);
	int tests;cin>>tests;
	while(tests--){
		int n;cin>>n;
		vector<vi> adj(n);
		rep(i,0,n-1){
			int a,b;cin>>a>>b;a--;b--;
			adj[a].pb(b);
			adj[b].pb(a);
		}
		if(n==2){
			cout<<"0\n";
			continue;
		}
		int start=0;
		while(sz(adj[start])==1) start++;
		vector<pair<pii,pii>> ops;
		map<pair<int,bool>,int> dp1;
		map<pair<int,bool>,bool> dp2;
		function<int(int,int,bool)> dfs1 = [&](int loc, int dad, bool up) -> int {
			if(sz(adj[loc])==1) return 0;
			// if(sz(adj[loc])<=2 && dad!=-1 && up==false) return 1+dfs1(loc,dad,true);
			if(dp1.count({loc,up})) return dp1[{loc,up}];
			int ans = 0;
			vector<pii> difs;
			for(int kid : adj[loc]) if(kid!=dad){
				int yes = dfs1(kid,loc,true);
				int no = dfs1(kid,loc,false);
				difs.pb({yes-no,kid});
				ans += no;
			}
			if(up || sz(difs)==1){
				assert(sz(difs)>0);
				auto itr = min_element(all(difs));
				ans += itr->F;
				ans += sz(difs)-1;
				for(auto p : difs){
					if(p.S == itr->S) dp2[{p.S,up}]=true;
					else dp2[{p.S,up}]=false;
				}

			}else{
				assert(sz(difs)>1);
				sort(all(difs));
				ans += difs[0].F + difs[1].F + sz(difs)-2;
				dp2[{difs[0].S,up}]=true;
				dp2[{difs[1].S,up}]=true;
				rep(i,2,sz(difs)) dp2[{difs[i].S,up}]=false;
			}
			return dp1[{loc,up}]=ans;
		};
		int res = dfs1(start,-1,false);
		cout<<res<<'\n';
		int used = 0;
		function<pii(int,int,bool)> dfs2 = [&](int loc, int dad, bool up) -> pii {
			// if(sz(adj[loc])<=2 && dad!=-1 && up==false) 
			pii me = {loc,loc};
			for(int kid : adj[loc]) if(kid!=dad){
				int ku = dp2[{kid,up}];
				if(!ku) continue;
				pii seg = dfs2(kid,loc,ku);
				if(me.S == loc) {
					assert(seg.F == kid);
					me.S = seg.S;
				}else{
					assert(me.F == loc);
					assert(seg.F==kid);
					me.F=seg.S;
				}
			}
			for(int kid : adj[loc]) if(kid!=dad){
				int ku = dp2[{kid,up}];
				if(ku) continue;
				pii seg = dfs2(kid,loc,ku);
				used++;
				cout<<loc+1<<' '<<kid+1<<' '<<seg.F+1<<' '<<me.S+1<<endl;
				me.S=seg.S;
			}
			if(up) {
				// if(me.F != loc){
				// 	cout<<loc<<' '<<dad<<' '<<up<<' '<<me.F<<' '<<me.S<<endl;
				// }
				assert(me.F == loc);
			}
			return me;
		};
		dfs2(start,-1,false);
		// cout<<res<<' '<<used<<endl;
		assert(res==used);
		
	}
}


Comments

Submit
0 Comments
More Questions

1118A - Water Buying
1462C - Unique Number
301A - Yaroslav and Sequence
38A - Army
38C - Blinds
1197A - DIY Wooden Ladder
1717D - Madoka and The Corruption Scheme
1296D - Fight with Monsters
729D - Sea Battle
788A - Functions again
1245B - Restricted RPS
1490D - Permutation Transformation
1087B - Div Times Mod
1213B - Bad Prices
1726B - Mainak and Interesting Sequence
1726D - Edge Split
1726C - Jatayu's Balanced Bracket Sequence
1726A - Mainak and Array
1613C - Poisoned Dagger
475B - Strongly Connected City
652B - z-sort
124B - Permutations
1496C - Diamond Miner
680B - Bear and Finding Criminals
1036E - Covered Points
1015D - Walking Between Houses
155B - Combination
1531A - Зингер | color
1678A - Tokitsukaze and All Zero Sequence
896A - Nephren gives a riddle